#!/usr/bin/env python3
# -*- encoding: utf8 -*-
import numpy as np
from enum import Enum
from skimage.draw import circle
from skimage.filters import gaussian
from skimage.transform import resize
from skimage.morphology import binary_closing, medial_axis, skeletonize, square

import warnings
warnings.filterwarnings("ignore", category=UserWarning)

def inverse_dist_transform(medial_axis):
    result = np.zeros(medial_axis.shape, dtype=np.uint8)
    pts = np.array(np.where(medial_axis > 0)).transpose()
    for pt in pts:
        dist = medial_axis[tuple(pt)]
        rr, cc = circle(pt[0], pt[1], dist, shape = medial_axis.shape)
        result[rr, cc] = 1
    return result

class ScaleFlag(str, Enum):
    thinning = 'thinning'
    fitting_target = 'fitting_target'

def morphology_scale(img, target_shape,
    mix = 0.2, width_adjustment = ScaleFlag.thinning,
    target_stroke_width = 20, thinning = 0.8,
    gaussian_sigma = 0, prune_dist_ratio = 0):
    '''Morphologically scaling a binary image

    Parameters
    ----------
    img: numpy.ndarray
        A binary image with 1 representing the foreground and 0 the background
    target_shape: tuple
        A tuple indicating the dimension of the target scaling shape
    mix: float
        Ranges from 0 to 1, the ratio to mix the direct scaling result and the
        morphological scaling result. 0 for fully morphological result
    width_adjustment: string
        Either 'thinning' or 'fitting_target'
    target_stroke_width: int
        Available when the 'fitting_target' flag is on. This argument gives the
        target stroke width to preserve in the scaled img
    thinning: float
        Available when the 'thinning' flag is on. Ratio to thin (or erode) the
        stroke width
    gaussian_sigma: float
        Gausssian sigma used for smoothing the distance field
    prune_dist_ratio: float
        The ratio of sizes of the pruned skeleton branches to the stroke width.

    Returns
    ----------
    scaled_img: numpy.ndarray
        The morphologically scaled binary image with 1 representing the
        foreground and 0 the background
    '''
    skeleton, dist = medial_axis(img, return_distance = True)
    stroke_width = 2 * np.average(dist[np.where(skeleton)])
    resized = resize(img, target_shape,
        anti_aliasing = None, mode = None)
    
    # options
    thins = width_adjustment == ScaleFlag.thinning
    fits_target = width_adjustment == ScaleFlag.fitting_target
    # the resized distance field
    resized_dist = gaussian(resize(dist, target_shape), sigma = gaussian_sigma)*\
        (target_stroke_width/stroke_width if fits_target else 1)
    _, dist_of_resized = medial_axis(resized, return_distance = True)
    # the dist used for scaling
    target_dist = (resized_dist * (1 - mix) + dist_of_resized * mix) *\
        (thinning if thins else 1)

    # skeleton of the resized 
    resized_skeleton = skeletonize(binary_closing(resized, square(4)))
    # prune skeleton
    pruned_skeleton = resized_skeleton &\
        ( target_dist >  stroke_width * prune_dist_ratio )

    return inverse_dist_transform(target_dist * pruned_skeleton)

